import os
import glob
import pandas as pd
import geopandas as gpd
import rasterio
import seaborn as sns

from src import utils, cdl, pvwater
import data_dir


if __name__ == '__main__':
    
    years = [2014, 2016]
    
    merged = []
    for year, file_name in data_dir.pvwater.items():
        if year not in years:
            continue

        # load PVWater
        df = pvwater.load(year=year, file_name=file_name)
        # find corresponding files for CDL
        raster_file, = glob.glob(os.path.join(
            data_dir.cdl_pajaro_folder,
            f'CDL_{year:d}_*.tif'))
        meta_file, = glob.glob(os.path.join(
            data_dir.cdl_pajaro_folder,
            f'CDL_{year:d}_*.tif.vat.dbf'))
        out_file = os.path.join(
            data_dir.cdl_vs_pvwater_folder,
            f'{year:d}.shp')

        # merge CDL with PVWater
        df = cdl.load(
            year=year,
            raster_file=raster_file,
            meta_file=meta_file,
            shape_file=df,
            out_file=out_file,
        )
        merged.append(df.loc[:, ['pvwater', 'cdl']])

    merged = pd.concat(merged)
    
    # SAVE CSV
    df = pd.crosstab(merged['pvwater'], merged['cdl'])
    df.to_csv(data_dir.cdl_vs_pvwater_csv)

    # PLOT FIGURE
    # plot the cross tabulation
    utils.plot_crosstab(
        input_df=merged, x_key='pvwater', y_key='cdl',
        x_title='PVWMA Land Use Type',
        y_title='USDA CDL Land Use Type',
        file_name=data_dir.cdl_vs_pvwater_fig,
        vmin=0, vmax=300,
        drop_threshold=20,
        drop_strs=['Ag. Unknown'],
        x_order=[
            'Strawberries',
            'Deciduous (Apple Orchards)',
            'Vegetable Row Crop',
            'Cover Crop',
            'Raspberries and Blackberries',
            'Blueberries',
            'Vines/Grapes',
            'Fallow',
            'Nurseries/Flowers/Tropical Plants',
            'Greenhouses',
            'Deciduous (Non-Apple Orchards)',
            'Natural',
            'Industrial',
            'Residential',
            'Other',
        ],
        y_order=[
            'Strawberries',
            'Shrubland',
            'Alfalfa',
            'Winter Wheat',
            'Mixed Forest',
            'Evergreen Forest',
            'Fallow/Idle Cropland',
            'Herbaceous Wetlands',
            'Grass/Pasture',
            'Developed/Open Space',
            'Developed/Low Intensity',
            'Developed/Med Intensity',
            'Developed/High Intensity',
            'Other',
        ],
    )
